# capture_requests.py
from mitmproxy import http
import json
import atexit
import sys
import os
import configparser
import datetime

# 默认配置
DEFAULT_CONFIG_FILE = "mitmproxy.conf"
TARGET_HOSTS = []
ALLOWED_HOSTS = []
BLOCKED_HOSTS = []
OUTPUT_FILE = "captured_requests.json"
SSL_INSECURE = True


def load_config(config_file=DEFAULT_CONFIG_FILE):
    """从配置文件加载参数"""
    global TARGET_HOSTS, ALLOWED_HOSTS, BLOCKED_HOSTS, OUTPUT_FILE, SSL_INSECURE

    if not os.path.exists(config_file):
        print(f"⚠️ 配置文件 {config_file} 不存在，使用默认配置")
        return

    try:
        config = configparser.ConfigParser()
        config.read(config_file, encoding='utf-8')

        # 读取设置
        if 'settings' in config:
            settings = config['settings']
            if 'target_hosts' in settings:
                TARGET_HOSTS = [host.strip() for host in settings['target_hosts'].split(',')]
            if 'output_file' in settings:
                OUTPUT_FILE = settings['output_file']

        # 读取SSL设置
        if 'ssl' in config:
            ssl_section = config['ssl']
            if 'ssl_insecure' in ssl_section:
                SSL_INSECURE = ssl_section.getboolean('ssl_insecure')

        # 读取过滤规则
        if 'filter' in config:
            filter_section = config['filter']
            if 'allowed_hosts' in filter_section:
                ALLOWED_HOSTS = [host.strip() for host in filter_section['allowed_hosts'].split(',')]
            if 'blocked_hosts' in filter_section:
                BLOCKED_HOSTS = [host.strip() for host in filter_section['blocked_hosts'].split(',')]

        print(f"✅ 成功加载配置文件: {config_file}")
        if TARGET_HOSTS:
            print(f"🎯 目标hosts: {TARGET_HOSTS}")
        if ALLOWED_HOSTS:
            print(f"✅ 允许的hosts: {ALLOWED_HOSTS}")
        if BLOCKED_HOSTS:
            print(f"🚫 阻止的hosts: {BLOCKED_HOSTS}")
        print(f"💾 输出文件: {OUTPUT_FILE}")

    except Exception as e:
        print(f"❌ 读取配置文件时出错: {e}")


# 在模块加载时自动加载配置
load_config()


def should_intercept(host):
    """判断是否应该拦截指定host的请求"""
    # 首先检查是否在阻止列表中
    if host in BLOCKED_HOSTS:
        return False

    # 检查通配符阻止列表
    for blocked_host in BLOCKED_HOSTS:
        if blocked_host.startswith('*') and host.endswith(blocked_host[1:]):
            return False

    # 如果配置了允许列表，则只允许列表中的host
    if ALLOWED_HOSTS:
        if host in ALLOWED_HOSTS:
            return True
        # 检查通配符允许列表
        for allowed_host in ALLOWED_HOSTS:
            if allowed_host.startswith('*') and host.endswith(allowed_host[1:]):
                return True
        return False

    # 如果配置了特定目标host，则只拦截这些host
    if TARGET_HOSTS:
        if host in TARGET_HOSTS:
            return True
        # 检查通配符匹配
        for target_host in TARGET_HOSTS:
            if target_host.startswith('*') and host.endswith(target_host[1:]):
                return True
        return False

    # 默认拦截所有请求
    return True


def response(flow: http.HTTPFlow) -> None:
    """捕获每个响应"""
    # 只拦截配置的host
    if not should_intercept(flow.request.host):
        return

    try:
        request_info = {
            "url": flow.request.pretty_url,
            "method": flow.request.method,
            "host": flow.request.host,
            "scheme": flow.request.scheme,  # HTTP or HTTPS
            "headers": dict(flow.request.headers),
            "status_code": flow.response.status_code,
            "response_headers": dict(flow.response.headers)
        }

        # 如果有请求体
        if flow.request.content:
            try:
                request_info["request_body"] = flow.request.content.decode('utf-8', errors='ignore')
            except:
                request_info["request_body"] = "[无法解码的请求体]"

        # 如果有响应体
        if flow.response.content:
            try:
                request_info["response_body"] = flow.response.content.decode('utf-8', errors='ignore')
            except:
                request_info["response_body"] = "[无法解码的响应体]"

        captured_requests.append(request_info)
        print(
            f"[{flow.response.status_code}] {flow.request.method} {flow.request.pretty_url} ({flow.request.scheme.upper()})")

    except Exception as e:
        print(f"处理请求时出错: {e}")


def save_requests():
    """保存请求到文件"""
    if captured_requests:
        try:
            # 生成带时间戳的文件名
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            name_part, ext_part = OUTPUT_FILE.rsplit('.', 1) if '.' in OUTPUT_FILE else (OUTPUT_FILE, 'json')
            final_file = f"{name_part}_{timestamp}.{ext_part}"

            with open(final_file, 'w', encoding='utf-8') as f:
                json.dump(captured_requests, f, indent=2, ensure_ascii=False)
            print(f"✅ 已保存 {len(captured_requests)} 个请求到 {final_file}")
        except Exception as e:
            print(f"保存文件时出错: {e}")
    else:
        print("没有捕获到任何请求")


def get_requests():
    """获取捕获的请求列表"""
    return captured_requests.copy()


def clear_requests():
    """清空捕获的请求列表"""
    captured_requests.clear()


# 全局变量存储请求
captured_requests = []

# 注册退出处理函数
atexit.register(save_requests)

# 提供addons列表供新版本mitmproxy使用
addons = []

# 初始化时显示配置信息
if TARGET_HOSTS:
    print(f"🎯 只拦截以下host的请求: {TARGET_HOSTS}")
elif ALLOWED_HOSTS:
    print(f"✅ 只允许以下host的请求: {ALLOWED_HOSTS}")
else:
    print("🌐 拦截所有host的请求（未配置特定host）")

if BLOCKED_HOSTS:
    print(f"🚫 阻止以下host的请求: {BLOCKED_HOSTS}")

print(f"🔐 SSL不安全模式: {'开启' if SSL_INSECURE else '关闭'}")
